package main

import (
	"context"
	"fmt"
	"golang.org/x/sync/semaphore"
	"sync"
	"sync/atomic"
)

type TaskPool struct {
	Ctx        context.Context
	TaskCap    int
	TaskChan   chan func()
	Concurrent *semaphore.Weighted
	State      int64 // 0 close 1 open
	FinishChan chan bool
	wg         *sync.WaitGroup
}

func NewTaskPool(ctx context.Context, cap int, gNum int) (*TaskPool, error) {
	if cap <= 0 {
		return nil, fmt.Errorf("init task cap no legal")
	}
	if gNum <= 0 {
		return nil, fmt.Errorf("init task groutine num no legal")
	}
	p := &TaskPool{
		Ctx:        ctx,
		TaskCap:    cap,
		TaskChan:   make(chan func(), cap),
		Concurrent: semaphore.NewWeighted(int64(gNum)),
		State:      1,
		FinishChan: make(chan bool, 1),
		wg:         &sync.WaitGroup{},
	}

	go func() {
		for {
			select {
			case tmpTask, ok := <-p.TaskChan:
				if !ok {
					p.wg.Wait()
					p.FinishChan <- true
					return
				} else {
					go p.Run(tmpTask)
				}

			}
		}
	}()
	return p, nil
}

func (p *TaskPool) Run(fn func()) {
	p.wg.Add(1)
	_ = p.Concurrent.Acquire(p.Ctx, 1)
	defer func() {
		p.wg.Done()
		p.Concurrent.Release(1)
	}()
	fn()
}

func (p *TaskPool) Close() {
	if !atomic.CompareAndSwapInt64(&p.State, 1, 0) {
		return
	}
	close(p.TaskChan)
}

func (p *TaskPool) Put(fn func()) error {
	if !atomic.CompareAndSwapInt64(&p.State, 1, 1) {
		return fmt.Errorf("task chan is close")
	}
	p.TaskChan <- fn
	return nil
}

func main() {
	var total int64 = 0

	pool, err := NewTaskPool(context.Background(), 100, 100)
	if err != nil {
		fmt.Println(err)
	}

	for i := 0; i < 100000; i++ {
		_ = pool.Put(func() {
			atomic.AddInt64(&total, 1)
		})
	}
	pool.Close()

	select {
	case <-pool.FinishChan:
		fmt.Println(total)
	}
}
